import os, tqdm, json, pickle, argparse
import numpy as np
from scipy.spatial.transform import Rotation as R

def label_converter(gt_path, pred_path, eval_name):
    if not os.path.exists(os.path.join(eval_name, 'gt_labels')):
        os.makedirs(os.path.join(eval_name, 'gt_labels'), exist_ok=True)
    
    if not os.path.exists(os.path.join(eval_name, 'pred_labels')):
        os.makedirs(os.path.join(eval_name, 'pred_labels'), exist_ok=True)

    # convert gt annotations
    with open(gt_path,'rb') as p:
        samples = pickle.load(p)
        
    for sample in tqdm.tqdm(samples, desc='Converting gt Labels...'):
        label_list = []
        sample_token = sample['sample_token']
        scene_token = sample['scene_token']
        timestamp = sample['timestamp']
        
        annos = sample['ann_infos']


        for anno in annos:
            name = anno['category_name']

            w, l, h = anno['size']
            x, y, z = anno['translation']

            rotation = anno['rotation']
            r = R.from_quat(rotation)
            yaw, pitch, roll = r.as_euler('zyx')

            label_list.append(f'{name.lower()} {x} {y} {z} {l} {w} {h} {yaw} \n')
            # label_list.append([name,x,y,z,l,w,h,yaw,'\n'])

        with open(f'{eval_name}/gt_labels/{sample_token}.csv', 'w') as gt_file:
            gt_file.writelines(label_list)

    # convert pred annotations
    with open(pred_path, 'r') as f:
        pred_labes = json.load(f)
    for sample in tqdm.tqdm(pred_labes['results'].keys(), desc='Converting pred Labels...'):
        sample_token = sample
        pred_list = []
        for label in pred_labes['results'][sample]:
            x1,y1,z1 = label['translation']
            w1,l1,h1 = label['size']
            rotation1 = label['rotation']
            rotation1 = [rotation1[1],rotation1[2],rotation1[3],rotation1[0]]
            roll1,pitch1,yaw1 = R.from_quat(rotation1).as_euler('xyz')
            score = label['detection_score']
            name1 = label['detection_name']
            pred_list.append(f'{name1.lower()} {x1} {y1} {z1} {l1} {w1} {h1} {yaw1} {score} \n')
            # pred_list.append([name1,x1,y1,z1,l1,w1,h1,yaw1,score,'\n'])
        with open(f'{eval_name}/pred_labels/{sample_token}.csv', 'w') as pred_file:
            pred_file.writelines(pred_list)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gt_path', type=str, required=True, help="Data path to gt annotations")
    parser.add_argument('--pred_path', type=str, required=True, help="Data path to pred annotations")
    parser.add_argument('--eval_name', type=str, required=True, help="eval exp name")
    args = parser.parse_args()
    
    gt_path = args.gt_path
    pred_path = args.pred_path
    eval_name = args.eval_name
    label_converter(gt_path, pred_path, eval_name)





